# -*- coding:utf-8 -*-

import os
from sklearn import model_selection
from sklearn.model_selection import train_test_split
import numpy as np
import shutil
from tqdm import tqdm, trange

def check_dir(file_path):
    if os.path.exists(file_path) ==False:
        os.makedirs(file_path)
    # else:
    #     # print("输出文件夹非空")
    #     shutil.rmtree(file_path)
    #     os.mkdir(file_path)

'''
根据交叉验证比例进行划分，可能造成类别不均衡
'''
def split_data_by_shuffleSplit(dataX, dataY, n_splits_use=10, test_size_use=0.2,random_state_use=0):
    sss = model_selection.ShuffleSplit(n_splits=n_splits_use, test_size=test_size_use, random_state=random_state_use)
    for train_index, test_index in sss.split(dataY):
    # print("TRAIN:", train_index, "TEST:", test_index)
        trainX, testX = dataX[train_index], dataX[test_index]
        trainY, testY = dataY[train_index], dataY[test_index]
    return trainX, trainY, testX, testY
 

def split_data_by_image_id(img_dir_path):
    path=os.listdir(img_dir_path)
    # print('类别',path)
    ii=0
    trainDataX_path_list = []
    trainDataY_path_list = []
    testDataX_path_list = []
    testDataY_path_list = []
    for category_calss in path:
        img_path_list = os.listdir(img_dir_path +  category_calss)
        for img_file_name in img_path_list:
            img_path = img_dir_path +  category_calss + os.sep + img_file_name
            # print('图片：',img_path)
            image_no = int(img_file_name.split('.png')[0])
            if image_no % 2 == 0:
                testDataX_path_list.append(img_path)
                testDataY_path_list.append(category_calss)
            else:
                trainDataX_path_list.append(img_path)
                trainDataY_path_list.append(category_calss)
    trainX, trainY = np.asarray(trainDataX_path_list), np.asarray(trainDataY_path_list)
    testX, testY = np.asarray(testDataX_path_list), np.asarray(testDataY_path_list)
    return trainX, trainY, testX, testY
   

def obtain_image_database(img_dir_path):
    '''
    img_dir_path ='images/'
    '''
    img_path_list = []
    img_file_path_list = []
    dataX_path_list = []
    dataY_path_list = []
    
    categories=os.listdir(img_dir_path)
    print('类别',categories)
    ii=0
    for category_calss in categories:
        img_path_list = os.listdir(img_dir_path +  category_calss)
        for img_file_name in img_path_list:
            img_path = img_dir_path +  category_calss + os.sep + img_file_name
            # print('图片：',img_path)
            dataX_path_list.append(img_path)
            dataY_path_list.append(category_calss)
    with open('label.txt', encoding='utf-8', mode='w+') as writer:
        count = 0
        for category_calss in categories:
            writer.write(str(count) + '\t' + category_calss + '\n')
            count = count + 1
    writer.close()
    return np.asarray(dataX_path_list), np.asarray(dataY_path_list), categories    
    

def copy_files_by_mode( dataX_use, dataY_use, classes, stage='train', img_root_dir_path='.'):
    with open(img_root_dir_path + os.sep + stage + '.txt', encoding='utf-8', mode='w+') as writer:
        for i in tqdm(range(len(dataX_use))):
            item_x = dataX_use[i]
            item_y = dataY_use[i]
            # print(item_x,item_y)           
            # 父级文件夹
            parent_dir_name = item_y
            
            total_save_path = img_root_dir_path + os.sep + stage
            
            final_file_save_dir = os.path.join(total_save_path,parent_dir_name)
            check_dir(final_file_save_dir)
            # print('保存文件完整路径：%s\n' % final_file_save)
            file_path = item_x
            filename = item_x.split(item_y + os.sep)[1] 
            
            final_file_save = os.path.join(final_file_save_dir,filename)
            shutil.copy(file_path,final_file_save)
            index_class_img = classes.index(item_y)
            writer.write(item_x + '\t' +  str(index_class_img)  + '\n')
    writer.close()
       
            
if __name__ == '__main__':
    img_dir_path = 'IMAGE.DYNA_FULL/'
    img_save_dir_path = "DYNA_FULL"
    dataX, dataY,categories = obtain_image_database(img_dir_path)
    print(dataX.shape)
    print(dataY.shape)
    
    # 策略1——根据样本比例分配训练测试集
    trainX,trainY,testX,testY = split_data_by_shuffleSplit(dataX, dataY, 10, 0.2, 0)


    # 策略二——根据图片id分配训练测试集
    # trainX,trainY,testX,testY = split_data_by_image_id(img_dir_path)
    
    use_stage = 'train'
    copy_files_by_mode(trainX,trainY,classes=categories, stage=use_stage, img_root_dir_path=img_save_dir_path)
    use_stage = 'test'
    copy_files_by_mode(testX,testY,classes=categories, stage=use_stage, img_root_dir_path=img_save_dir_path)
    
        